package query

import (
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"sync"
)

// HygonQuery Hygon GPU查询实现
type HygonQuery struct{}

func init() {
	RegisterQueryer("hygon", func() GPUQuery { return &HygonQuery{} })
}

func (q *HygonQuery) Query() ([]GPUInfo, error) {
	// 获取设备列表
	devices, err := q.getDevices()
	if err != nil {
		return nil, fmt.Errorf("获取设备列表失败: %v", err)
	}

	// 如果设备数量为0，直接返回空结果
	if len(devices) == 0 {
		return []GPUInfo{}, nil
	}

	// 获取所有设备的使用率信息
	usageOutput, err := runCommand("hy-smi", "--showuse")
	if err != nil {
		return nil, fmt.Errorf("获取使用率信息失败: %v", err)
	}

	// 获取所有设备的内存信息
	memOutput, err := runCommand("hy-smi", "--showmeminfo", "vram")
	if err != nil {
		return nil, fmt.Errorf("获取内存信息失败: %v", err)
	}

	// 并行解析设备信息
	var wg sync.WaitGroup
	gpuInfos := make([]GPUInfo, len(devices))
	errors := make([]error, len(devices))

	for i, device := range devices {
		wg.Add(1)
		go func(index int, device int) {
			defer wg.Done()
			utilization, err := q.parseUtilization(usageOutput, device)
			if err != nil {
				errors[index] = err
				return
			}

			memoryUsed, memoryTotal, err := q.parseMemoryInfo(memOutput, device)
			if err != nil {
				errors[index] = err
				return
			}

			gpuInfos[index] = GPUInfo{
				Index:       device,
				Utilization: utilization,
				MemoryUsed:  memoryUsed,
				MemoryTotal: memoryTotal,
			}
		}(i, device)
	}
	wg.Wait()

	// 检查是否有错误
	for _, err := range errors {
		if err != nil {
			return nil, fmt.Errorf("解析设备信息失败: %v", err)
		}
	}

	return gpuInfos, nil
}

// getDevices 获取设备列表
func (q *HygonQuery) getDevices() ([]int, error) {
	output, err := runCommand("hy-smi", "--showid")
	if err != nil {
		return nil, err
	}

	// 使用正则表达式快速提取设备ID
	deviceRegex := regexp.MustCompile(`DCU\[(\d+)\]`)
	lines := strings.Split(output, "\n")
	var devices []int

	for _, line := range lines {
		if matches := deviceRegex.FindStringSubmatch(line); matches != nil {
			if deviceID, err := strconv.Atoi(matches[1]); err == nil {
				devices = append(devices, deviceID)
			}
		}
	}

	if len(devices) == 0 {
		return nil, fmt.Errorf("未找到有效的GPU设备")
	}

	return devices, nil
}

// parseUtilization 解析使用率信息
func (q *HygonQuery) parseUtilization(output string, device int) (int, error) {
	// 预编译正则表达式
	utilRegex := regexp.MustCompile(`DCU\[` + strconv.Itoa(device) + `\]\s+: DCU use \(%\): (\d+)`)
	matches := utilRegex.FindStringSubmatch(output)
	if matches == nil {
		return 0, fmt.Errorf("未找到设备%d的使用率信息", device)
	}

	if utilization, err := strconv.Atoi(matches[1]); err == nil {
		return utilization, nil
	}

	return 0, fmt.Errorf("解析使用率失败")
}

// parseMemoryInfo 解析内存信息
func (q *HygonQuery) parseMemoryInfo(output string, device int) (int, int, error) {
	// 预编译正则表达式
	memTotalRegex := regexp.MustCompile(`DCU\[` + strconv.Itoa(device) + `\]\s+: vram Total Memory \(MiB\): (\d+)`)
	memUsedRegex := regexp.MustCompile(`DCU\[` + strconv.Itoa(device) + `\]\s+: vram Total Used Memory \(MiB\): (\d+)`)

	// 查找内存总量
	totalMatches := memTotalRegex.FindStringSubmatch(output)
	if totalMatches == nil {
		return 0, 0, fmt.Errorf("未找到设备%d的内存总量信息", device)
	}
	total, err := strconv.Atoi(totalMatches[1])
	if err != nil {
		return 0, 0, fmt.Errorf("解析内存总量失败: %v", err)
	}

	// 查找内存使用量
	usedMatches := memUsedRegex.FindStringSubmatch(output)
	if usedMatches == nil {
		return 0, 0, fmt.Errorf("未找到设备%d的内存使用量信息", device)
	}
	used, err := strconv.Atoi(usedMatches[1])
	if err != nil {
		return 0, 0, fmt.Errorf("解析内存使用量失败: %v", err)
	}

	return used, total, nil
}
